{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "228bb2a6",
   "metadata": {},
   "source": [
    "# Adding new model to the existing DataLabeler pipeline"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cab7a569",
   "metadata": {},
   "source": [
    "Consider the case when we would like to explore different character-level neural network models and evaluate their performance on different datasets. The existing DataLabeler in the DataProfiler library already contains a preprocessor, a postprocessor, and a character-level CNN (Convolutional Neural Network) model that are combined to work on such data. All we need is to build additional model classes that inherit the main functionalities from the CNN model and also adapt the model construction to the desired architectures. In this example, we define such a new model to be used with the Data Labeler component of the Data Profiler. In particular, a character-level LSTM (Long Short-Term Memory) model is implemented, then integrated into the DataLabeler pipeline to be trained with a tabular dataset. The process includes the following steps:\n",
    "\n",
    "    - Build a new character-level LSTM model that inherits the CNN model\n",
    "    - Load the DataLabeler from the DataProfiler\n",
    "    - Swap the existing CNN model with the new LSTM model\n",
    "    - Train the data labeler pipeline on a given dataset\n",
    "\n",
    "First, let's import the libraries needed for this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16624c48",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import json\n",
    "import pandas as pd\n",
    "sys.path.insert(0, '..')\n",
    "import dataprofiler as dp"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e90728ab",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3d61981c",
   "metadata": {},
   "source": [
    "In this example, we use a structured dataset, the aws honeypot dataset, given in the test folder of the library. This dataset is first read by the Data Reader class of the Data Profiler, then split into training and test data to be used in the next sections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f031fe06",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use data reader to read input data\n",
    "data = dp.Data(\"../dataprofiler/tests/data/csv/aws_honeypot_marx_geo.csv\")\n",
    "df_data = data.data\n",
    "\n",
    "# split data to training and test set\n",
    "split_ratio = 0.2\n",
    "df_data = df_data.sample(frac=1).reset_index(drop=True)\n",
    "data_train = df_data[:int((1 - split_ratio) * len(df_data))]\n",
    "data_test = df_data[int((1 - split_ratio) * len(df_data)):]\n",
    "\n",
    "df_data.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "745ed0d4",
   "metadata": {},
   "source": [
    "## Implement a new character-level LSTM model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7375b0c0",
   "metadata": {},
   "source": [
    "This new model is inherited from `CharacterLevelCnnModel` class, with some modifications on the following functions\n",
    "\n",
    "`__init__`: to add new parameters for the LSTM model. The new parameters, `size_lstm`, `rec_dropout`, `activation`, `recurrent_activation`, specify number of LSTM layers, activation function, and recurrent dropout ratio.\n",
    "\n",
    "`_validate_parameters`: to add additional checks on the new parameters for the LSTM model\n",
    "\n",
    "`_construct_model`: to construct the new LSTM model with the desired architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8568fb49",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from dataprofiler.labelers.character_level_cnn_model import (\n",
    "    CharacterLevelCnnModel,\n",
    "    create_glove_char,\n",
    "    build_embd_dictionary,\n",
    ")\n",
    "from dataprofiler.labelers.base_model import BaseModel\n",
    "from dataprofiler.labelers.labeler_utils import F1Score\n",
    "\n",
    "\n",
    "# CharacterLevelLstmModel derives from CharacterLevelCnnModel\n",
    "#########################################################\n",
    "#########################################################\n",
    "class CharacterLevelLstmModel(CharacterLevelCnnModel):\n",
    "    # boolean if the label mapping requires the mapping for index 0 reserved\n",
    "    requires_zero_mapping = True\n",
    "\n",
    "    def __init__(self, label_mapping=None, parameters=None):\n",
    "        \"\"\"\n",
    "        LSTM Model Initializer\n",
    "        \"\"\"\n",
    "\n",
    "        # parameter initialization\n",
    "        if not parameters:\n",
    "            parameters = {}\n",
    "        parameters.setdefault(\"max_length\", 3400)\n",
    "        parameters.setdefault(\"max_char_encoding_id\", 127)\n",
    "        parameters.setdefault(\"dim_embed\", 64)\n",
    "        parameters.setdefault(\"size_fc\", [32, 32])\n",
    "        parameters.setdefault(\"dropout\", 0.1)\n",
    "        # new parameters for LSTM model\n",
    "        #########################################################\n",
    "        #########################################################\n",
    "        parameters.setdefault(\"size_lstm\", [64])\n",
    "        parameters.setdefault(\"rec_dropout\", 0.1)\n",
    "        parameters.setdefault(\"activation\", \"tanh\")\n",
    "        parameters.setdefault(\"recurrent_activation\", \"sigmoid\")\n",
    "        #########################################################\n",
    "        #########################################################\n",
    "        parameters.setdefault(\"default_label\", \"UNKNOWN\")\n",
    "        parameters[\"pad_label\"] = \"PAD\"\n",
    "        self._epoch_id = 0\n",
    "\n",
    "        # reconstruct flags for model\n",
    "        self._model_num_labels = 0\n",
    "        self._model_default_ind = -1\n",
    "\n",
    "        BaseModel.__init__(self, label_mapping, parameters)\n",
    "\n",
    "    def _validate_parameters(self, parameters):\n",
    "        \"\"\"\n",
    "        Validate the parameters sent in. Raise error if invalid parameters are\n",
    "        present.\n",
    "        \"\"\"\n",
    "        errors = []\n",
    "        list_of_necessary_params = [\n",
    "            \"max_length\",\n",
    "            \"max_char_encoding_id\",\n",
    "            \"dim_embed\",\n",
    "            \"size_fc\",\n",
    "            \"dropout\",\n",
    "            \"size_lstm\",\n",
    "            \"rec_dropout\",\n",
    "            \"activation\",\n",
    "            \"recurrent_activation\",\n",
    "            \"default_label\",\n",
    "            \"pad_label\",\n",
    "        ]\n",
    "        # Make sure the necessary parameters are present and valid.\n",
    "        for param in parameters:\n",
    "            if param in [\n",
    "                \"max_length\",\n",
    "                \"max_char_encoding_id\",\n",
    "                \"dim_embed\",\n",
    "                \"size_conv\",\n",
    "            ]:\n",
    "                if (\n",
    "                    not isinstance(parameters[param], (int, float))\n",
    "                    or parameters[param] < 0\n",
    "                ):\n",
    "                    errors.append(\n",
    "                        param + \" must be a valid integer or float \" \"greater than 0.\"\n",
    "                    )\n",
    "            elif param in [\n",
    "                \"dropout\",\n",
    "                \"rec_dropout\",\n",
    "            ]:  # additional check for rec_dropout\n",
    "                if (\n",
    "                    not isinstance(parameters[param], (int, float))\n",
    "                    or parameters[param] < 0\n",
    "                    or parameters[param] > 1\n",
    "                ):\n",
    "                    errors.append(\n",
    "                        param + \" must be a valid integer or float \" \"from 0 to 1.\"\n",
    "                    )\n",
    "            elif (\n",
    "                param == \"size_fc\" or param == \"size_lstm\"\n",
    "            ):  # additional check for size_lstm\n",
    "                if (\n",
    "                    not isinstance(parameters[param], list)\n",
    "                    or len(parameters[param]) == 0\n",
    "                ):\n",
    "                    errors.append(param + \" must be a non-empty list of \" \"integers.\")\n",
    "                else:\n",
    "                    for item in parameters[param]:\n",
    "                        if not isinstance(item, int):\n",
    "                            errors.append(\n",
    "                                param + \" must be a non-empty \" \"list of integers.\"\n",
    "                            )\n",
    "                            break\n",
    "            elif param in [\n",
    "                \"default_label\",\n",
    "                \"activation\",\n",
    "                \"recurrent_activation\",\n",
    "            ]:  # additional check for activation and recurrent_activation\n",
    "                if not isinstance(parameters[param], str):\n",
    "                    error = str(param) + \" must be a string.\"\n",
    "                    errors.append(error)\n",
    "\n",
    "        # Error if there are extra parameters thrown in\n",
    "        for param in parameters:\n",
    "            if param not in list_of_necessary_params:\n",
    "                errors.append(param + \" is not an accepted parameter.\")\n",
    "        if errors:\n",
    "            raise ValueError(\"\\n\".join(errors))\n",
    "\n",
    "    def _construct_model(self):\n",
    "        \"\"\"\n",
    "        Model constructor for the data labeler. This also serves as a weight\n",
    "        reset.\n",
    "\n",
    "        :return: None\n",
    "        \"\"\"\n",
    "        num_labels = self.num_labels\n",
    "        default_ind = self.label_mapping[self._parameters[\"default_label\"]]\n",
    "\n",
    "        # Reset model\n",
    "        tf.keras.backend.clear_session()\n",
    "\n",
    "        # generate glove embedding\n",
    "        create_glove_char(self._parameters[\"dim_embed\"])\n",
    "\n",
    "        # generate model\n",
    "        self._model = tf.keras.models.Sequential()\n",
    "\n",
    "        # default parameters\n",
    "        max_length = self._parameters[\"max_length\"]\n",
    "        max_char_encoding_id = self._parameters[\"max_char_encoding_id\"]\n",
    "\n",
    "        # Encoding layer\n",
    "        def encoding_function(input_str):\n",
    "            char_in_vector = CharacterLevelLstmModel._char_encoding_layer(\n",
    "                input_str, max_char_encoding_id, max_length\n",
    "            )\n",
    "            return char_in_vector\n",
    "\n",
    "        self._model.add(tf.keras.layers.Input(shape=(None,), dtype=tf.string))\n",
    "\n",
    "        self._model.add(\n",
    "            tf.keras.layers.Lambda(encoding_function, output_shape=tuple([max_length]))\n",
    "        )\n",
    "\n",
    "        # Create a pre-trained weight matrix\n",
    "        # character encoding indices range from 0 to max_char_encoding_id,\n",
    "        # we add one extra index for out-of-vocabulary character\n",
    "        embed_file = os.path.join(\n",
    "            \"../dataprofiler/labelers\",\n",
    "            \"embeddings/glove-reduced-{}D.txt\".format(self._parameters[\"dim_embed\"]),\n",
    "        )\n",
    "        embedding_matrix = np.zeros(\n",
    "            (max_char_encoding_id + 2, self._parameters[\"dim_embed\"])\n",
    "        )\n",
    "        embedding_dict = build_embd_dictionary(embed_file)\n",
    "\n",
    "        input_shape = tuple([max_length])\n",
    "        # Fill in the weight matrix: let pad and space be 0s\n",
    "        for ascii_num in range(max_char_encoding_id):\n",
    "            if chr(ascii_num) in embedding_dict:\n",
    "                embedding_matrix[ascii_num + 1] = embedding_dict[chr(ascii_num)]\n",
    "\n",
    "        self._model.add(\n",
    "            tf.keras.layers.Embedding(\n",
    "                max_char_encoding_id + 2,\n",
    "                self._parameters[\"dim_embed\"],\n",
    "                weights=[embedding_matrix],\n",
    "                input_length=input_shape[0],\n",
    "                trainable=True,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        # Add the lstm layers\n",
    "        #########################################################\n",
    "        #########################################################\n",
    "        for size in self._parameters[\"size_lstm\"]:\n",
    "            self._model.add(\n",
    "                tf.keras.layers.LSTM(\n",
    "                    units=size,\n",
    "                    recurrent_dropout=self._parameters[\"rec_dropout\"],\n",
    "                    activation=self._parameters[\"activation\"],\n",
    "                    recurrent_activation=self._parameters[\"recurrent_activation\"],\n",
    "                    return_sequences=True,\n",
    "                )\n",
    "            )\n",
    "            if self._parameters[\"dropout\"]:\n",
    "                self._model.add(tf.keras.layers.Dropout(self._parameters[\"dropout\"]))\n",
    "        #########################################################\n",
    "        #########################################################\n",
    "\n",
    "        # Add the fully connected layers\n",
    "        for size in self._parameters[\"size_fc\"]:\n",
    "            self._model.add(tf.keras.layers.Dense(units=size, activation=\"relu\"))\n",
    "            if self._parameters[\"dropout\"]:\n",
    "                self._model.add(tf.keras.layers.Dropout(self._parameters[\"dropout\"]))\n",
    "\n",
    "        # Add the final Softmax layer\n",
    "        self._model.add(tf.keras.layers.Dense(num_labels, activation=\"softmax\"))\n",
    "\n",
    "        # Output the model into a .pb file for TensorFlow\n",
    "        argmax_layer = tf.keras.backend.argmax(self._model.output)\n",
    "\n",
    "        # Create confidence layers\n",
    "        final_predicted_layer = CharacterLevelLstmModel._argmax_threshold_layer(\n",
    "            num_labels, threshold=0.0, default_ind=default_ind\n",
    "        )\n",
    "\n",
    "        argmax_outputs = self._model.outputs + [\n",
    "            argmax_layer,\n",
    "            final_predicted_layer(argmax_layer, self._model.output),\n",
    "        ]\n",
    "        self._model = tf.keras.Model(self._model.inputs, argmax_outputs)\n",
    "\n",
    "        # Compile the model\n",
    "        softmax_output_layer_name = self._model.outputs[0].name.split(\"/\")[0]\n",
    "        losses = {softmax_output_layer_name: \"categorical_crossentropy\"}\n",
    "\n",
    "        # use f1 score metric\n",
    "        f1_score_training = F1Score(num_classes=num_labels, average=\"micro\")\n",
    "        metrics = {softmax_output_layer_name: [\"acc\", f1_score_training]}\n",
    "\n",
    "        self._model.compile(loss=losses, optimizer=\"adam\", metrics=metrics)\n",
    "\n",
    "        self._epoch_id = 0\n",
    "        self._model_num_labels = num_labels\n",
    "        self._model_default_ind = default_ind"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d66bd25c",
   "metadata": {},
   "source": [
    "## Integrate the new LSTM model to the DataLabeler"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "479f407a",
   "metadata": {},
   "source": [
    "Once the LSTM model is built, it replaces the existing model in the DataLabeler pipeline, which is then trained on the given dataset. Note that, as the DataLabeler is trained on the above tabular dataset, its label mapping is updated by the list of column names in that dataset while training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb482ffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get labels from the given dataset\n",
    "value_label_df = data_train.reset_index(drop=True).melt()\n",
    "value_label_df.columns = [1, 0]  # labels=1, values=0 in that order\n",
    "value_label_df = value_label_df.astype(str)\n",
    "labels = value_label_df[1].unique().tolist()\n",
    "\n",
    "# create a new LSTM model\n",
    "# set default label (one of the column names) to the model\n",
    "model = CharacterLevelLstmModel(label_mapping=labels, parameters={'default_label': 'comment'})\n",
    "\n",
    "# add the new LSTM model to the data labeler\n",
    "data_labeler = dp.DataLabeler(labeler_type='structured', trainable=True)\n",
    "data_labeler.set_model(model)\n",
    "\n",
    "# set default label (one of the column names) to the preprocessor and postprocessor\n",
    "processor_params = {'default_label': 'comment'}\n",
    "data_labeler._preprocessor.set_params(**processor_params)\n",
    "data_labeler._postprocessor.set_params(**processor_params)\n",
    "\n",
    "# train the data labeler\n",
    "save_dirpath=\"data_labeler_saved\"\n",
    "if not os.path.exists(save_dirpath):\n",
    "    os.makedirs(save_dirpath)\n",
    "\n",
    "epochs=2\n",
    "data_labeler.fit(\n",
    "    x=value_label_df[0], y=value_label_df[1], labels=labels, epochs=epochs)\n",
    "if save_dirpath:\n",
    "    data_labeler.save_to_disk(save_dirpath)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "14b78c69",
   "metadata": {},
   "source": [
    "The trained Data Labeler is then used by the Data Profiler to provide the prediction on the new dataset. In this example, all options except data labeler are disabled for the sake of presenting data labeler functionality. The results are given in the columnar format where true column types are given in the first column, and the predicted column labels are given in the second column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdfcf1d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# predict with the data labeler object\n",
    "profile_options = dp.ProfilerOptions()\n",
    "profile_options.set({\"structured_options.text.is_enabled\": False, \n",
    "                     \"int.is_enabled\": False, \n",
    "                     \"float.is_enabled\": False, \n",
    "                     \"order.is_enabled\": False, \n",
    "                     \"category.is_enabled\": False, \n",
    "                     \"datetime.is_enabled\": False,})\n",
    "profile_options.set({'structured_options.data_labeler.data_labeler_object': data_labeler})\n",
    "profile = dp.Profiler(data_test, options=profile_options)\n",
    "\n",
    "# get the prediction from the data profiler\n",
    "def get_structured_results(results):\n",
    "    columns = []\n",
    "    predictions = []\n",
    "    for col_report in results['data_stats']:\n",
    "        columns.append(col_report['column_name'])\n",
    "        predictions.append(col_report['data_label'])\n",
    "\n",
    "    df_results = pd.DataFrame({'Column': columns, 'Prediction': predictions})\n",
    "    return df_results\n",
    "\n",
    "results = profile.report()\n",
    "print(get_structured_results(results))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cc60ff8a",
   "metadata": {},
   "source": [
    "In summary, users can define their own model, plug it in the DataLabeler pipeline, and train the labeler with the new dataset. Above, we show one example of adding the LSTM model to the pipeline. Interested users can implement other neural network models as desired with the same process."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
